#include "stdafx.h"
#include "injector.h"
#include <stdint.h>

#include <XProcess>
#include <conio.h>

using namespace GenericFilter;
using namespace std;

void usage(const char *exe)
{
    cout << "Usage : " << endl;
    cout << exe << " {-i pid|-p process} [-undo]" << endl;
    cout << exe << " -i pid" << endl;
    cout << exe << " -p process name" << endl;
    exit(0);
}

enum class PorcessTarget
{
    Pid,
    Name
};

class Arguments
{
public:
    PorcessTarget target;
    string name;
    uint64_t pid;
    bool unInject = false;
public:
    static Arguments parse(int argc, char **argv)
    {
        Arguments args;
        for (int i = 1; i < argc; i++)
        {
            auto argstr = string(argv[i]);
        	if (argstr == "-i")
        	{
                args.target = PorcessTarget::Pid;
                args.pid = stoull(argv[++i]);
        	}
            else if (argstr == "-n")
            {
                args.target = PorcessTarget::Name;
                args.name = argv[++i];
            }
            else if (argstr == "-undo")
            {
                args.unInject = true;
            }
            else
            {
                cout << "unknow argument : " << argstr << endl;
                usage(argv[0]);
            }
        }
        return args;
    }
};

int main(int argc, char **argv)
{
    auto args = Arguments::parse(argc, argv);

    if (!args.unInject)
    {
        // inject 
        if (args.target == PorcessTarget::Name)
        {
            cout << "Injecting all process : " << args.name << " ..." << endl;

            Injector::inject(args.name);
        }
        else
        {
            cout << "Injecting pid : " << args.pid << " ..." << endl;

            Injector::inject(args.pid);
        }
    }
    else
    {
        // un-inject
        if (args.target == PorcessTarget::Name)
        {
            cout << "Undoing Injection  for all process with name : " << args.name << " ..." << endl;

            Injector::undoInject(args.name);
        }
        else
        {
            cout << "Undoing Injection of process with pid : " << args.pid << " ..." << endl;

            Injector::undoInject(args.pid);
        }
    }
    return 0;
}